#pragma once

#include <ostream>
#include "utility.h"

#ifndef NULL
#define NULL 0
#endif

template <typename Comparable>
class AvlTree
{
public:
	AvlTree()
	{
		root = NULL;
	}

	AvlTree(const AvlTree & rhs)
	{
		root = clone(rhs.root);
	}

	~AvlTree()
	{
		makeEmpty();
	}

	const Comparable & findMin() const
	{
		return findMin(root);
	}

	const Comparable & findMax() const
	{
		return findMax(root);
	}

	bool contains(const Comparable & x) const
	{
		return contains(x, root);
	}

	bool isEmpty() const
	{
		return root == NULL;
	}

	void printTree(std::ostream& out = std::cout) const
	{
		if (isEmpty())
			out << "Empty tree." << std::endl;
		else
			printTree(root, out);
	}

	void makeEmpty()
	{
		makeEmpty(root);
	}

	void insert(const Comparable & x)
	{
		insert(x, root);
	}

	void remove(const Comparable & x)
	{
		remove(x, root);
	}

	const AvlTree & operator=(const AvlTree & rhs)
	{
		if (this != &rhs)
		{
			makeEmpty();
			root = clone(rhs.root);
		}
		return *this;
	}

private:
	struct AvlNode
	{
		Comparable element;
		AvlNode *left;
		AvlNode *right;
		int height;

		AvlNode(const Comparable& theElement, AvlNode *lt,
			   	AvlNode *rt, int h = 0) : element(theElement),
				left(lt), right(rt), height(h)
		{}
	};
	AvlNode *root;
	
	void insert(const Comparable& x, AvlNode* &t)
	{
		if (t == NULL)
			t = new AvlNode(x, NULL, NULL);
		else if (x < t->element)
		{
			insert(x, t->left);
			if (height(t->left) - height(t->right) == 2)
				if (x < t->left->element)
					rotateWithLeftChild(t);
				else
					doubleWithLeftChild(t);
		}
		else if (t->element < x)
		{
			insert(x, t->right);
			if (height(t->right) - height(t->left) == 2)
				if (t->right->element < x)
					rotateWithRightChild(t);
				else
					doubleWithRightChild(t);
		}
		else
			; // Duplicate; do nothing
		t->height = mymax(height(t->left), height(t->right)) + 1; 
	}

	void remove(const Comparable& x, AvlNode* &t)
	{
		if (t == NULL)
			return;
		if (x < t->element)
			remove(x, t->left);
		else if (t->element < x)
			remove(x, t->right);
		else if (t->left != NULL && t->right != NULL)
		{
			t->element = findMin(t->right)->element;
			remove(t->element, t->right);
		}
		else
		{
			AvlNode* oldNode = t;
			t = (t->left != NULL) ? t->left : t->right;
			delete oldNode;
		}
	}

	AvlNode* findMin(AvlNode *t) const
	{
		if (t == NULL)
			return NULL;
		if (t->left == NULL)
			return t;
		return findMin(t->left);
	}

	AvlNode* findMax(AvlNode *t) const
	{
		if (t != NULL)
			while (t->right != NULL)
				t = t->right;
		return t;
	}

	bool contains(const Comparable &x, AvlNode *t) const
	{
		if (t == NULL)
			return false;
		else if (x < t->element)
			return contains(x, t->left);
		else if (t->element < x)
			return contains(x, t->right);
		else
			return true;
	}

	void makeEmpty(AvlNode* &t)
	{
		if (t != NULL)
		{
			makeEmpty(t->left);
			makeEmpty(t->right);
			delete t;
		}
		t = NULL;
	}

	void printTree(AvlNode *t, std::ostream& out) const
	{
		if (t != NULL)
		{
			printTree(t->left, out);
			out << t->element << std::endl;
			printTree(t->right, out);
		}
	}

	AvlNode* clone(AvlNode *t) const
	{
		if (t == NULL)
			return NULL;

		return new AvlNode(t->element, clone(t->left), clone(t->right));
	}

	int height(AvlNode* t) const
	{
		return t == NULL ? -1 : t->height;
	}

	void rotateWithLeftChild(AvlNode* &k2)
	{
		AvlNode* k1 = k2->left;
		k2->left = k1->right;
		k1->right = k2;
		k2->height = mymax(height(k2->left), height(k2->right)) + 1;
		k1->height = mymax(height(k1->left), k2->height) + 1;
		k2 = k1;
	}

	void rotateWithRightChild(AvlNode* &k2)
	{
		AvlNode* k1 = k2->right;
		k2->right = k1->left;
		k1->left = k2;
		k2->height = mymax(height(k2->left), height(k2->right)) + 1;
		k1->height = mymax(k2->height, height(k1->right)) + 1;
		k2 = k1;
	}

	void doubleWithLeftChild(AvlNode* &k3)
	{
		rotateWithRightChild(k3->left);
		rotateWithLeftChild(k3);
	}

	void doubleWithRightChild(AvlNode* &k3)
	{
		rotateWithLeftChild(k3->right);
		rotateWithRightChild(k3);
	}
};
